import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, LabelEncoder
import joblib
from collections import defaultdict
import clean_data


def train_do():
    """
    处理训练数据：
      - 连续变量：等频分箱 → WOE 编码（用于逻辑回归）
      - 类别变量：Label Encoding（用于 XGB/RF/LR）
      - 逻辑回归输入：仅 WOE 特征 + 类别编码，然后标准化
    同时保存所有转换器，供测试集使用。
    """
    data = pd.read_csv('../data/train.csv')
    y = data['Attrition'].copy()
    x = data[list(clean_data.Cleaning_data().keys())].copy()

    # === Step 1: WOE 编码（保存 bins 和 woe_map）===
    woe_features = ['Age', 'MonthlyIncome', 'TotalWorkingYears', 'DistanceFromHome']
    woe_transforms = {}  # 用于保存每个特征的 bins 和 woe_map

    for col in woe_features:
        # 使用 qcut 获取分箱边界（bins）
        _, bins = pd.qcut(x[col], q=5, duplicates='drop', retbins=True)

        # 用 bins 固定分箱（确保可复现）
        bin_col = f"{col}_bin"
        x[bin_col] = pd.cut(x[col], bins=bins, include_lowest=True)

        # 计算 WOE
        ct = pd.crosstab(x[bin_col], y)
        ct = ct + 0.5  # Laplace smoothing
        good_total, bad_total = ct[0].sum(), ct[1].sum()

        woe_map = {}
        for interval in ct.index:
            woe = np.log((ct.loc[interval, 0] / good_total) / (ct.loc[interval, 1] / bad_total))
            woe_map[interval] = woe

        # 应用 WOE
        x[f"{col}_WOE"] = x[bin_col].map(woe_map)

        # 清理临时列
        x.drop(columns=[bin_col], inplace=True)

        # 保存转换信息
        woe_transforms[col] = {
            'bins': bins,
            'woe_map': woe_map
        }

    # === Step 2: Label Encoding（保存每个 encoder）===
    cat_cols = [
        'BusinessTravel', 'Education', 'JobRole',
        'MaritalStatus', 'OverTime', 'EducationField', 'Department'
    ]
    label_encoders = {}

    for col in cat_cols:
        le = LabelEncoder()
        x[col] = le.fit_transform(x[col])
        label_encoders[col] = le  # 保存 encoder

    # === Step 4: 保存所有转换器 ===
    joblib.dump(woe_transforms, '../model/transform/woe_transforms.pkl')
    joblib.dump(label_encoders, '../model/transform/label_encoders.pkl')

    print("已保存：WOE 转换规则、LabelEncoders、StandardScaler")

    # 返回完整的 x（含原始连续列 + WOE 列 + 编码类别列），可用于 XGB/RF
    return x, y

